#ifndef QUADTREE_H
#define QUADTREE_H

/**
 * \bibitem{mark2008computational}
 * M. de Berg, O. Cheong, M. van Kreveld, and M. Overmars, \textit{Computational Geometry: Algorithms and Applications}, Springer, 2008.
 */

#include "vertex2d.h"
#include "box2d.h"
#include "table.h"
#include "key.h"
#include "cell.h"
#include <malloc.h>
#include <stack>
#include <queue>
#include <memory>

// Position code of leaf nodes
#define QUADTREE_LOWER_LEFT     0b00
#define QUADTREE_UPPER_LEFT     0b01
#define QUADTREE_LOWER_RIGHT    0b10
#define QUADTREE_UPPER_RIGHT    0b11

namespace RealScan
{

// QuadTree
template <typename T, unsigned BucketSize = 4>
class QuadTree : public Table<T>
{
public:
    // Constructor
    QuadTree(const Vertex2D& center, double radius) 
    {
        // Initialize root node
        m_root = std::make_unique<Box2D<T> >(center, radius);
    }

    // Destructor
    ~QuadTree() { malloc_trim(0); }

    ////// `Table` interface implementation //////

    // Search
    std::optional<T> search(const Key& key) const override
    {
        // Get the root node of the quadtree
        const Box2D<T>* root = m_root.get();

        // Traverse the tree from the root, moving through the internal (non-leaf) nodes
        while (!root->leaf)
        {
            // Determine the child node index based on the key's position
            unsigned ids = belong(Vertex2D(key.x, key.y), root);

            // If the child node exists, move to the child node
            if (root->child[ids])
            {
                root = root->child[ids].get();
            }
            else
            {
                // If the child node does not exist, the key is not found in the quadtree
                return std::nullopt;
            }
        }

        // At this point, we have reached a leaf node, so search the hash table for the key
        auto it = root->hashTable.find(Vertex2D(key.x, key.y));

        // If the key is found in the hash table, return the associated value
        if (it != root->hashTable.end())
        {
            return it->second;
        }

        // If the key is not found, return std::nullopt to indicate the key doesn't exist
        return std::nullopt;
    }

    // Insert
    bool insert(const Key& key, const T& value) override
    {
        insert(Vertex2D(key.x, key.y), value, m_root.get());
        return true;
    }

    // Remove
    void remove(const Key& key) override
    {
        // Stack to store the nodes that are traversed during the search
        std::stack<Box2D<T>*> stackedBoxes;
        stackedBoxes.push(m_root.get()); // Start with the root node

        Box2D<T>* box = stackedBoxes.top(); // Initialize the current box as the root

        // Traverse down the tree to find the leaf node containing the key
        while (!box->leaf)
        {
            // Determine which child node the key belongs to
            unsigned ids = belong(Vertex2D(key.x, key.y), box);

            // If the child node exists, move to that child and add it to the stack
            if (box->child[ids])
            {
                stackedBoxes.push(box->child[ids].get());
                box = stackedBoxes.top(); // Update the current box to the child node
            }
            else
            {
                // If the child node does not exist, return (key not found)
                return;
            }
        }

        // Once we are at a leaf node, search for the key in the hash table
        auto it = box->hashTable.find(Vertex2D(key.x, key.y));

        // If the key is found in the hash table, delete it
        if (it != box->hashTable.end())
        {
            box->hashTable.erase(it); // Remove the key-value pair from the hash table

            // After removal, attempt to clean up redundant boxes (maintain tree structure consistency)
            reduce(stackedBoxes); // Call the reduce function to potentially free unnecessary nodes
        }
    }

    // // Vectorize
    void toVector(std::vector<Cell<T>>& vec, const Key& minXY, const Key& maxXY) const override
    {
        // Clear the existing content of the vector
        vec.clear();

        // Use a queue to simulate the recursive process (BFS traversal)
        std::queue<Box2D<T>*> queueBoxes;
        queueBoxes.push(m_root.get());  // Start with the root node

        // Continue processing while the queue is not empty
        while (!queueBoxes.empty())
        {
            const Box2D<T>* front = queueBoxes.front();  // Get the node at the front of the queue

            // If the current node is a leaf node, process its elements
            if (front->leaf)
            {
                // Get the enclosure status of the current box with respect to the specified region
                EnclosureStatus status = getEnclosureStatus(front, Vertex2D(minXY.x, minXY.y), Vertex2D(maxXY.x, maxXY.y));

                switch (status)
                {
                    case CONTAINNING:  // The current box is completely contained within the specified region
                        // Add all elements in the hash table to the vector
                        for (auto pair : front->hashTable)
                        {
                            vec.push_back(Cell<T>(Key(pair.first.x, pair.first.y), pair.second));
                        }
                        break;

                    case INTERSECTING:  // The current box intersects with the specified region
                        // Add elements that are within the region
                        for (auto pair : front->hashTable)
                        {
                            if (vertexInRegion(pair.first, Vertex2D(minXY.x, minXY.y), Vertex2D(maxXY.x, maxXY.y)))
                            {
                                vec.push_back(Cell<T>(Key(pair.first.x, pair.first.y), pair.second));
                            }
                        }
                        break;

                    case DISJOINTING:  // The current box does not intersect with the specified region
                        // Do nothing
                        break;
                }
            }
            else  // If the current node is not a leaf node, traverse its children
            {
                // Check each child node (4 children for a quadtree)
                for (int i = 0; i < 4; ++i)
                {
                    if (front->child[i])
                    {
                        // Get the enclosure status of the child node with respect to the specified region
                        EnclosureStatus status = getEnclosureStatus(front->child[i].get(), Vertex2D(minXY.x, minXY.y), Vertex2D(maxXY.x, maxXY.y));

                        switch (status)
                        {
                            case CONTAINNING:  // The child node is completely contained within the region
                                // Add all vertices from this child node to the vector
                                addAllVertexToVector(front->child[i].get(), vec);
                                break;

                            case INTERSECTING:  // The child node intersects with the region
                                // Add the child node to the queue for further processing
                                queueBoxes.push(front->child[i].get());
                                break;

                            case DISJOINTING:  // The child node does not intersect with the region
                                // Do nothing
                                break;
                        }
                    }
                }
            }

            // Pop the front node from the queue after processing
            queueBoxes.pop();
        }
    }

    // cells number
    std::size_t size() const override
    {
        std::size_t size_ = 0;
        size(m_root.get(), size_);
        return size_;
    }

    // memory used, const there
    std::size_t space() const override
    {
        std::size_t bytes = 0;
        space(m_root.get(), bytes);
        return bytes; 
    }

private:
    // root node
    std::unique_ptr<Box2D<T> > m_root; 

    // Determine which child node the vertex potentially belongs to in the quadtree
    unsigned belong(const Vertex2D& vertex, const Box2D<T>* root) const
    {
        unsigned X = 0, Y = 0;
        X |= ((vertex.x >= root->center.x) << 1);
        Y |= ((vertex.y >= root->center.y) << 0);
        return (X | Y);
    }

    // Get the center position of the specified child node of the root
    Vertex2D getCenter(unsigned ids, const Box2D<T>* root) const
    {
        // Offset from the root center to the child center
        Vertex2D delta(0, 0);

        // Calculate the offset based on the child node's ID
        switch (ids)
        {
            case QUADTREE_LOWER_LEFT:
                // Lower-left quadrant: offset both x and y negatively
                delta.x = -root->radius / 2.0;
                delta.y = -root->radius / 2.0;
                break;

            case QUADTREE_UPPER_LEFT:
                // Upper-left quadrant: offset x negatively and y positively
                delta.x = -root->radius / 2.0;
                delta.y =  root->radius / 2.0;
                break;

            case QUADTREE_LOWER_RIGHT:
                // Lower-right quadrant: offset x positively and y negatively
                delta.x =  root->radius / 2.0;
                delta.y = -root->radius / 2.0;
                break;

            case QUADTREE_UPPER_RIGHT:
                // Upper-right quadrant: offset both x and y positively
                delta.x =  root->radius / 2.0;
                delta.y =  root->radius / 2.0;
                break;
        }

        // Add the offset to the root center to get the child node's center position
        return (root->center + delta);
    }

    // Retrieve the child node of the root that may contain the given vertex.
    // If the child node does not exist, create it.
    Box2D<T>* getChild(const Vertex2D& vertex, Box2D<T>* root)
    {
        // Determine which child node the vertex belongs to based on its coordinates.
        unsigned ids = belong(vertex, root);

        // If the child node does not exist, create it.
        if (!root->child[ids])
        {
            // Calculate the radius of the new child node (half the root's radius).
            double newRadius = root->radius / 2.0;

            // Create a new child node at the calculated position with the new radius.
            // The position is determined by the center of the child node.
            root->child[ids] = std::make_unique<Box2D<T>>(getCenter(ids, root), newRadius);
        }

        // Return a pointer to the child node.
        return root->child[ids].get();
    }

        // Attempt to insert a (vertex, value) pair into the root node
    void insert(const Vertex2D& vertex, const T& value, Box2D<T>* root)
    {
        // If the current node is a leaf node, try to insert into its bucket
        if (root->leaf)
        {
            // If the vertex already exists in the bucket, update its value
            if (root->hashTable.find(vertex) != root->hashTable.end())
            {
                root->hashTable[vertex] = value;
            }
            // If the vertex does not exist, attempt to insert it
            else
            {
                // If the bucket has sufficient capacity, directly insert the vertex and value
                if (root->hashTable.size() < BucketSize)
                {
                    root->hashTable[vertex] = value;
                }
                // If the bucket is full, split the leaf node into child nodes
                else
                {
                    // Convert the leaf node to a branch node
                    root->leaf = false;

                    // Distribute the existing bucket elements to the appropriate child nodes
                    for (const auto& pair : root->hashTable)
                    {
                        insert(pair.first, pair.second, getChild(pair.first, root));
                    }

                    // Clear the bucket after moving its elements
                    root->hashTable.clear();

                    // Insert the new (vertex, value) pair into the appropriate child node
                    insert(vertex, value, getChild(vertex, root));
                }
            }
        }
        // If the current node is a branch node, recursively insert into the appropriate child node
        else
        {
            insert(vertex, value, getChild(vertex, root));
        }
    }

    // After removing a node, check whether the parent Box should reclaim all its child Boxes
    void reduce(std::stack<Box2D<T>*>& stackedBoxes)
    {
        // Initially assume that the parent node can be reduced (reclaimed)
        bool canReduce = true;

        // Pop the current node to move back to the parent node
        stackedBoxes.pop();

        // Continue checking parent nodes as long as reduction is possible and the stack is not empty
        while (canReduce && !stackedBoxes.empty())
        {
            // Get the parent node from the stack
            Box2D<T>* parent = stackedBoxes.top();

            // Count the total number of elements across all child nodes of the parent
            std::size_t numKeys = 0;
            for (int i = 0; i < 4; ++i)
            {
                if (parent->child[i])
                {
                    // If the child is a leaf node, add its bucket size to the count
                    if (parent->child[i]->leaf)
                    {
                        numKeys += parent->child[i]->hashTable.size();
                    }
                    // If the child is not a leaf node, reduction is not possible
                    else return;
                }
            }

            // Determine if the total number of keys across all children is within the allowable limit
            canReduce &= (numKeys <= BucketSize);

            // If reduction is possible, reclaim all child nodes
            if (canReduce)
            {
                for (int i = 0; i < 4; ++i)
                {
                    if (parent->child[i])
                    {
                        // Transfer all elements from the child's bucket to the parent's bucket
                        for (const auto& pair : parent->child[i]->hashTable)
                        {
                            parent->hashTable[pair.first] = pair.second;
                        }
                        // Reset the child node (deallocate memory)
                        parent->child[i].reset();
                    }
                }
                // Mark the parent as a leaf node after reclaiming its children
                parent->leaf = true;
            }

            // Pop the current parent node and move up the stack to the next parent
            stackedBoxes.pop();
        }
    }

    // Check if a vertex is within a specified region
    bool vertexInRegion(const Vertex2D& vertex, const Vertex2D& minXY, const Vertex2D& maxXY) const
    {
        // Return true if the vertex lies within the region defined by minXY and maxXY (inclusive)
        // This checks that:
        // 1. The vertex's coordinates are greater than or equal to the minimum coordinates (minXY)
        // 2. The vertex's coordinates are less than or equal to the maximum coordinates (maxXY)
        return ((minXY <= vertex) && (vertex <= maxXY));
    }

    // Enumeration to represent the enclosure status between two regions
    enum EnclosureStatus
    {
        DISJOINTING,    // The two regions are disjoint, with no overlap
        INTERSECTING,   // The two regions intersect but neither fully contains the other
        CONTAINNING     // One region completely contains the other
    };

    // Determine the relationship between a box and a given rectangular region
    EnclosureStatus getEnclosureStatus(const Box2D<T>* box, const Vertex2D& minXY, const Vertex2D& maxXY) const
    {
        // Initialize the counter for the number of box corners enclosed by the region
        std::size_t enclosedPts = 0;

        // Check if each corner of the box is inside the region (minXY, maxXY)
        enclosedPts += vertexInRegion({box->center.x - box->radius, box->center.y - box->radius}, minXY, maxXY); // Bottom-left corner
        enclosedPts += vertexInRegion({box->center.x - box->radius, box->center.y + box->radius}, minXY, maxXY); // Top-left corner
        enclosedPts += vertexInRegion({box->center.x + box->radius, box->center.y - box->radius}, minXY, maxXY); // Bottom-right corner
        enclosedPts += vertexInRegion({box->center.x + box->radius, box->center.y + box->radius}, minXY, maxXY); // Top-right corner

        // If all 4 corners of the box are enclosed by the region, the box is fully contained within the region
        if (enclosedPts == 4)
        {
            return CONTAINNING;
        }
        // If at least one corner is inside the region, the box and the region are intersecting
        else if (enclosedPts > 0)
        {
            return INTERSECTING;
        }
        else
        {
            // If none of the box corners are inside the region, check if the region is inside the box
            // Define the box's bounding box (minXY and maxXY as the corners of the box)
            Vertex2D boxMin(box->center.x - box->radius, box->center.y - box->radius);
            Vertex2D boxMax(box->center.x + box->radius, box->center.y + box->radius);

            // Check if each corner of the region (minXY, maxXY) is inside the box's bounding box
            enclosedPts += vertexInRegion(minXY, boxMin, boxMax); // minXY inside the box
            enclosedPts += vertexInRegion(maxXY, boxMin, boxMax); // maxXY inside the box
            enclosedPts += vertexInRegion({minXY.x, maxXY.y}, boxMin, boxMax); // top-left of the region inside the box
            enclosedPts += vertexInRegion({maxXY.x, minXY.y}, boxMin, boxMax); // bottom-right of the region inside the box

            // If at least one corner of the region is inside the box, they are intersecting
            if (enclosedPts > 0)
            {
                return INTERSECTING;
            }
        }

        // If no corners of the region are inside the box, the two are disjoint (no overlap)
        return DISJOINTING;
    }

    // Recursively add all vertices from the box and its descendants to the vector
    void addAllVertexToVector(const Box2D<T>* box, std::vector<Cell<T> >& vec) const
    {
        // If the current box is a leaf node, add all elements (vertices) in the box's bucket to the vector
        if (box->leaf)
        {
            // Iterate through all items in the box's bucket and add them to the vector
            for (auto pair : box->hashTable)
            {
                // Create a Grid object using the vertex's coordinates and value, and push it into the vector
                vec.push_back(Cell<T>(Key(pair.first.x, pair.first.y), pair.second));
            }
        }
        // If the current box is not a leaf node, recursively add elements from its child nodes
        else
        {
            // Iterate through all 4 child nodes of the current box
            for (int i = 0; i < 4; ++i)
            {
                // If the child node exists, recursively add its elements to the vector
                if (box->child[i])
                {
                    addAllVertexToVector(box->child[i].get(), vec);
                }
            }
        }
    }

    // Recursively calculate the memory usage of the quadtree
    void space(const Box2D<T>* root, std::size_t& bytes) const
    {   
        // If the box pointer is null, return immediately
        if (!root) return;

        // If the current node is a leaf, add the space occupied by its elements (the bucket)
        if (root->leaf)
        {
            // object size
            bytes += sizeof(root->hashTable);

            // buckets size
            bytes += root->hashTable.bucket_count() * sizeof(void*);

            // elements size
            for (std::size_t i = 0; i < root->hashTable.bucket_count(); ++i)
            {
                bytes += root->hashTable.bucket_size(i) * (sizeof(T) + sizeof(Vertex2D) + sizeof(void*));
            }
        }

        // Add the space occupied by the node itself (Box<T>)
        bytes += (sizeof(Vertex2D) + 4 * sizeof(void*));

        // Recursively calculate the space for each child node
        for (int i = 0; i < 4; ++i)
        {
            // If the child node exists, calculate its space recursively
            space(root->child[i].get(), bytes);
        }
    }

    // Recursively calculate the total number of elements in the quadtree
    void size(const Box2D<T>* root, std::size_t& size_) const
    {
        // If the box pointer is null, return immediately
        if (!root) return;

        // If the current node is a leaf, add the number of elements in its bucket
        if (root->leaf)
        {
            size_ += root->hashTable.size(); // Add the number of elements in the bucket
        }
        else
        {
            // If the current node is not a leaf, recursively calculate the size for each child node
            for (int i = 0; i < 4; ++i)
            {
                // Recursively calculate the size for each child
                size(root->child[i].get(), size_);
            }
        }
    }

};

} /* RealScan */

#endif /* QUADTREE_H */